import os
import numpy as np
import pandas as pd
import cv2
from pathlib import Path
import shutil  
from time import time

def get_mean_and_std(img):
    x_mean, x_std = cv2.meanStdDev(img)
    x_mean = np.hstack(x_mean).astype("float16")
    x_std = np.hstack(x_std).astype("float16")
    return x_mean, x_std

def rectify(img_path, code, out_dir, mean_std):
    mean_avg, std_avg = mean_std
    
    sc = cv2.imread(img_path)
    sc = cv2.cvtColor(sc, cv2.COLOR_BGR2LAB).astype("float32")
    s_mean, s_std = get_mean_and_std(sc)

    sc = sc.astype("float16")
    img_n = ((sc-s_mean)*(std_avg/s_std)) + mean_avg
    np.putmask(img_n, img_n > 255, 255)
    np.putmask(img_n, img_n < 0, 0)
    dst = cv2.cvtColor(cv2.convertScaleAbs(img_n.astype('float32')), cv2.COLOR_LAB2BGR)

    img_name = str(Path(img_path).name)
    dst_path = os.path.join(out_dir, code)
    os.makedirs(dst_path, exist_ok=True)
    cv2.imwrite(os.path.join(dst_path, img_name), dst)
    return os.path.join(dst_path, img_name)


def color_transfer(df, img_col_name, code_col_name, out_dir, mean_std):
    if os.path.exists(out_dir):
        shutil.rmtree(out_dir)
        os.makedirs(out_dir)

    img_col_name_new = img_col_name + '_new'

    try:
        from pandarallel import pandarallel
        pandarallel.initialize(progress_bar=True) 
        print('Use multi threading !')
        is_pandarallel = True
    except:
        print('Use single threading !')
        is_pandarallel = False

    start_time = time()
    if is_pandarallel:
        df[img_col_name_new] = df.parallel_apply(lambda x:rectify(img_path=x[img_col_name], code=x[code_col_name], out_dir=out_dir, mean_std=mean_std), axis=1)
    else:
        df[img_col_name_new] = df.apply(lambda x:rectify(img_path=x[img_col_name], code=x[code_col_name], out_dir=out_dir, mean_std=mean_std), axis=1)

    df = df.drop([img_col_name], axis=1)
    df = df.rename(columns={img_col_name_new: img_col_name})
    print('\n color transfer (%d imgs) cost %.2f s.' % (len(df), time()-start_time))
    return df

